#!/usr/bin/python
#coding:utf-8
'''
卷积神经网络-MobileNet实例，深度可分类卷积,多分类器
'''
import tensorflow as tf 
import os 
import _pickle as cPickle 
import numpy as np 

CIFAR_DIR="./cifar-10-batches-py"

def load_data(filename):
    '''
    从文件读取数据
    '''
    with open(filename,'rb') as f:
        data=cPickle.load(f,encoding='bytes')
        return data[b'data'],data[b'labels']

class CifarData:
    '''
    数据处理类
    '''
    def __init__(self,fileList,need_shuffle): #need_shuffle 是否需要shuffle
        all_data=[]
        all_labelList = []
        for filename in fileList:
            data,lableList= load_data(filename)
            
            #获取全量十个类数据
            all_data.append(data)
            all_labelList.append(lableList)

            #获取两类数据
            # for item,label in zip(data,lableList):
            #     if label in [0,1]:
            #         all_data.append(item)
            #         all_labelList.append(label)
        
        self._data= np.vstack(all_data)  #纵行合并为矩阵
        self._data = self._data /127.5 -1 #数据缩放 增加归一化
        self._labels= np.hstack(all_labelList) #横向合并为一维向量
        self._num_example=self._data.shape[0]
        self._need_shuffle = need_shuffle
        self._indicator = 0

        if self._need_shuffle:
            self._shuffle_data()
    
    def _shuffle_data(self):
        '''
        打乱数据
        '''
        p = np.random.permutation(self._num_example)  #实现数据混排
        self._data=self._data[p]
        self._labels=self._labels[p]
    
    def nex_batche(self,batch_size):
        '''
        每次返回batch_size个样本
        数据获取一遍后，shuffle一遍，重新获取
        '''
        end_indicator= self._indicator + batch_size
        if end_indicator > self._num_example:
            if self._need_shuffle:
                self._shuffle_data()
                self._indicator=0
                end_indicator = batch_size
            else:
                raise Exception('have no more examples')
        
        if end_indicator > self._num_example:
            raise Exception("batch size is larger than all example")

        batch_data=self._data[self._indicator:end_indicator]
        batch_labels = self._labels[self._indicator: end_indicator]
        self._indicator=end_indicator

        return batch_data,batch_labels

train_filenames=[os.path.join(CIFAR_DIR,'data_batch_%d'% i) for i in range(1,6)]
test_filenames =[os.path.join(CIFAR_DIR,'test_batch')]

train_data=CifarData(train_filenames,True)
test_data=CifarData(test_filenames,False)


def separable_conv_block(x,output_channel_number,name):
    '''
    mobile：可分离卷积
    args:
        x:输入数据
        output_channel_number:输出通道数，output channel of 1x1 conv layer
        name:名字
    '''
    with tf.variable_scope(name):
        input_channel = x.get_shape().as_list()[-1]
        #channel_wise_x:[channel1,channel2,...]
        channel_wise_x = tf.split(x,input_channel,axis=3)
        output_channels = []
        for i in range(len(channel_wise_x)):
            output_channel= tf.layers.conv2d(channel_wise_x[i],
                                            1,
                                            [3,3],
                                            strides=(1,1),
                                            padding='same',
                                            activation=tf.nn.relu,
                                            name='conv_%d' %i)
            output_channels.append(output_channel)
        concat_layer = tf.concat(output_channels,axis=3)
        conv1_1 = tf.layers.conv2d(concat_layer,
                                    output_channel_number,
                                    (1,1),
                                    strides=(1,1),
                                    padding='same',
                                    activation=tf.nn.relu,
                                    name = 'conv1_1')
        
    return conv1_1


#构建计算图
x= tf.placeholder(tf.float32,[None,3072]) #placeholder(占位符
y= tf.placeholder(tf.int64,[None]) #None 样本数不确定

#数据转换
x_image=tf.reshape(x,[-1,3,32,32])   # 格式转换
x_image=tf.transpose(x_image,perm=[0,2,3,1])  #通道转换，编程正确顺序

#第一层 卷积层 conv1：神经元图 feature_map 输出图像
conv1= tf.layers.conv2d(x_image,
                        32,  #output channel number 输出通道数
                        (3,3), #kernel size卷积核大小
                        padding='same', #padding类型，same（做padding）、valid（不做padding）
                        activation = tf.nn.relu, #激活函数
                        name= 'conv1')

#第一层 池化层 大小减小一半 16*16
pooling1 =tf.layers.max_pooling2d(conv1,
                        (2,2), #kernel size 内核大小
                        (2,2), #stride 步长
                        name='pool1')

separable_2a=separable_conv_block(pooling1,
                            32,
                            name='separable_2a')
separable_2b=separable_conv_block(separable_2a,
                            32,
                            name='separable_2b')

#第一层 池化层 大小减小一半 16*16
pooling2 =tf.layers.max_pooling2d(separable_2b,
                        (2,2), #kernel size 内核大小
                        (2,2), #stride 步长
                        name='pool2')


separable_3a=separable_conv_block(pooling2,
                            32,
                            name='separable_3a')
separable_3b=separable_conv_block(separable_3a,
                            32,
                            name='separable_3b')

#第一层 池化层 大小减小一半 16*16
pooling3 =tf.layers.max_pooling2d(separable_3b,
                        (2,2), #kernel size 内核大小
                        (2,2), #stride 步长
                        name='pool3')

#展开
flatten= tf.layers.flatten(pooling3)

#全连接层
y_ = tf.layers.dense(flatten,10)

###########平方差损失函数###########
#针对当前数据集，该损失函验证集成功率40%
#api: e^x / sum(e^x)
# p_y= tf.nn.softmax(y_)
# #5 --> [0000100000]
# y_one_hot= tf.one_hot(y,10,dtype=tf.float32)
# loss=tf.reduce_mean(tf.square(y_one_hot - p_y))  #reduce_mean 求均值  square：求平方，结果为真实值和预测值差的平方

###########交叉熵损失函数##########
#针对当前数据集，损失函数验证集成功率30%
loss=tf.losses.sparse_softmax_cross_entropy(labels=y,logits=y_)

'''
#sigmoid 
#[None,10]
p_y_1 = tf.nn.sigmoid(y_)  #传入sigmoid得到概率值
#[None,1]
y_reshaped=tf.reshape(y,(-1,1))  #数据类型 int64
y_reshaped_float = tf.cast(y_reshaped,tf.float32)  #类型转换 int64 转 float32

#平方差计算损失函数
loss=tf.reduce_mean(tf.square(y_reshaped_float - p_y_1))  #reduce_mean 求均值  square：求平方，结果为真实值和预测值差的平方
'''

#准确率
predict =tf.argmax(y_,1) 

#[1,0,1,1...]
correct_prediction= tf.equal(predict,y) #得到正确预测值
accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float64)) #得到准确率

#梯度下降方法
with tf.name_scope('train_op'):
    train_op = tf.train.AdamOptimizer(1e-3).minimize(loss)  #1e-3：10的负三次方

#创建绘画
init = tf.global_variables_initializer()

batch_size=20
train_steps=100000
test_steps=100

with tf.Session() as sess:
    sess.run(init)
    for i in range(train_steps):
        batch_data,batch_labels=train_data.nex_batche(batch_size)
        loss_val,acc_val,_=sess.run(
                    [loss,accuracy,train_op],
                    feed_dict={x:batch_data, y:batch_labels})  #加train_op 表示训练，不加表示计算，feed_dict:输入数据
        
        if i % 500 ==0:
            print("[Train] Step: %d, loss: %4.5f, acc: %4.5f" %(i,loss_val,acc_val))
        
        #每5000次test数据集评测
        if (i+1) % 5000 ==0:
            test_data=CifarData(test_filenames,False)
            all_test_acc_val=[]
            for j in range(test_steps):
                test_batch_data,test_batch_labels=test_data.nex_batche(batch_size)
                test_acc_val=sess.run(
                    [accuracy],
                    feed_dict={
                        x:test_batch_data,
                        y:test_batch_labels
                    })
                all_test_acc_val.append(test_acc_val)
            
            test_acc=np.mean(all_test_acc_val)
            print('[Test ] Step: %d, acc: %4.5f' % (i+1,test_acc))